'''
flask API
initialize yolo_pose_model
api1       get_pose_by_yolo_batch 参数中的yolo_model直接使用上面初始化的
api2       draw_pose 
'''
from demo_yolo_pose import get_yolo_pose,get_pose_by_yolo_batch
from demo_yolo_pose import draw_pose
from util_for_argparse import get_cuda_port_args

from flask import Flask, request, jsonify, send_file
import numpy as np
import cv2
import io

args = get_cuda_port_args()

app = Flask(__name__)

# 初始化 yolo_pose_model
yolo_pose_model = get_yolo_pose()

@app.route('/api/get_pose_by_yolo_batch', methods=['POST'])
def api_get_pose_by_yolo_batch():
    """
    输入：支持图片文件上传 或 图片url列表
    返回：每张图片的关键点results（直接转成列表）
    """
    # 处理图片数据
    imgs = []
    img_paths = []

    if 'images' in request.files:
        files = request.files.getlist('images')
        for f in files:
            file_bytes = np.frombuffer(f.read(), np.uint8)
            img = cv2.imdecode(file_bytes, cv2.IMREAD_COLOR)
            if img is None:
                return jsonify({'error': 'Cannot decode uploaded image'}), 400
            imgs.append(img)
    elif request.is_json:
        data = request.get_json()
        img_paths = data.get('img_paths', [])
        if not img_paths:
            return jsonify({'error': 'No images found'}), 400
    else:
        return jsonify({'error': 'No valid input provided'}), 400

    # 调用 pose api
    try:
        if imgs:
            results = get_pose_by_yolo_batch(yolo_pose_model, imgs=imgs)
        elif img_paths:
            results = get_pose_by_yolo_batch(yolo_pose_model, img_paths=img_paths)
        else:
            return jsonify({'error': 'No images to process'}), 400

        # results 可以遍历，提取 .keypoints.data
        all_keypoints = []
        for result in results:
            kpts = result.keypoints.data.cpu().numpy().tolist()  # shape: [N, 17, 3]
            all_keypoints.append(kpts)
        return jsonify({'keypoints': all_keypoints})
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/api/draw_pose', methods=['POST'])
def api_draw_pose():
    """
    输入：
      - image 文件
      - kps: 关键点，json格式 [[x,y,score],...] 或 [[[x,y,score],...], ...] (支持一个或多个targets)
    返回：
      - 绘制后图片（jpg格式字节流）
    """
    if 'image' not in request.files or 'kps' not in request.form:
        return jsonify({'error': 'Missing image or kps'}), 400

    image_file = request.files['image']
    kps_str = request.form['kps']
    try:
        kps = np.array(eval(kps_str))
    except Exception as e:
        return jsonify({'error': f'Invalid kps: {e}'}), 400

    file_bytes = np.frombuffer(image_file.read(), np.uint8)
    img = cv2.imdecode(file_bytes, cv2.IMREAD_COLOR)
    if img is None:
        return jsonify({'error': 'Cannot decode uploaded image'}), 400

    # 可选参数
    threshold = float(request.form.get('threshold', 0.3))
    radius = int(request.form.get('radius', 4))

    # 绘制
    try:
        result_img = draw_pose(img_arr=img, kps=kps, threshold=threshold, radius=radius)
    except Exception as e:
        return jsonify({'error': str(e)}), 500

    # 返回图片（jpg二进制流）
    _, img_encoded = cv2.imencode('.jpg', result_img)
    return send_file(
        io.BytesIO(img_encoded.tobytes()),
        mimetype='image/jpeg',
        as_attachment=False,
        download_name='pose.jpg'
    )

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=args.port, debug=False)
